package org.example.myleet.rosalind.chbp;

import org.example.myleet.rosalind.utils.NewickTreeNode;
import org.example.myleet.rosalind.utils.NewickTreeUtil;

import java.io.*;
import java.util.*;

public class Solution {

    public static void main(String[] args) throws IOException {
        Solution solution = new Solution();
        String path = solution.getClass().getResource("").getPath() + "example.txt";
        try (BufferedReader br = new BufferedReader(new InputStreamReader(new FileInputStream(path)))) {
            String line = br.readLine().trim();
            String[] species = line.split(" ");
            List<String> taxaList = Arrays.asList(species);
            int n = species.length;
            List<RowSum> characterTable = new ArrayList<>();
            while (null != (line = br.readLine())) {
                line = line.trim();
                if (line.length() < 1) {
                    break;
                }
                characterTable.add(new RowSum(line, n));
            }
            //对characterTable每一行进行排序，按照每一行的熵值递减排序，使得1和0的数量越接近的row越排在前面
            characterTable.sort((o1, o2) -> {
                if (o1.entropy < o2.entropy) {
                    return 1;
                } else if (o1.entropy > o2.entropy) {
                    return -1;
                }
                return 0;
            });
            //物种对应的下标
            Map<String, Integer> taxaMapIndex = new HashMap<>();
            for (int i = 0; i < n; ++i) {
                taxaMapIndex.put(species[i], i);
            }
            //从根节点开始重新构建物种树（无根二叉树），这里给出的根节点是一个trivial节点
            NewickTreeNode root = new NewickTreeNode(0, 0, "", null);
            reconstructTree(root, characterTable, 0, taxaList, taxaMapIndex);
            //按Newick格式序列化二叉树输出结果
            StringBuilder res = new StringBuilder();
            NewickTreeUtil.serialize(root, res);
            res.append(';');
            System.out.println(res);
        }
    }

    /**
     * 递归中使用的rows，是按照熵值由大到小排序的，能够保证靠前的row是靠近树顶节点的划分，靠后的row是针对靠近叶子节点的划分
     */
    private static void reconstructTree(NewickTreeNode parent, List<RowSum> rows, int depth, List<String> taxaList, Map<String, Integer> taxaMapIndex) {
        if (depth >= rows.size()) {
            //构建树的深度超过给出的输入行的数量时，不再继续递归，检查taxaList的大小
            if (taxaList.size() == 1) {
                //如果大小是1，则相当于parent是nontrivial节点
                parent.setValue(taxaList.get(0));
            } else if (taxaList.size() == 2) {
                //如果大小是2，则相当于parent是trivial节点，下面挂2个叶子节点
                parent.getChildren().add(new NewickTreeNode(0, 0, taxaList.get(0), parent));
                parent.getChildren().add(new NewickTreeNode(0, 0, taxaList.get(1), parent));
            }
            return;
        }
        if (taxaList.size() == 1) {
            //如果大小是1，则相当于parent是nontrivial节点，此时无需继续深入分割
            parent.setValue(taxaList.get(0));
            return;
        }
        //下面是递归过程中需要处理的业务，每行输入的0、1组成的信息代表的是题目所给的物种应该按怎样的方式进行划分
        String row = rows.get(depth).row;
        //按0划分为左子树的节点，1划分为右子树的节点
        List<String> leftTaxaList = new ArrayList<>();
        List<String> rightTaxaList = new ArrayList<>();
        for (String taxa : taxaList) {
            int j = taxaMapIndex.get(taxa);
            if (row.charAt(j) == '0') {
                leftTaxaList.add(taxa);
            } else {
                rightTaxaList.add(taxa);
            }
        }
        if (leftTaxaList.size() < 1 || rightTaxaList.size() < 1) {
            //左子树或右子树为空的情况下，说明本次划分并没有得到有效节点，因此需要忽略本行的信息，选取下一行进行划分
            reconstructTree(parent, rows, depth + 1, taxaList, taxaMapIndex);
        } else {
            //能够划分为左右子树，说明找到一个trivial节点，继续深入递归划分各自的部分（分治）
            NewickTreeNode left = new NewickTreeNode(0, 0, "", parent);
            reconstructTree(left, rows, depth + 1, leftTaxaList, taxaMapIndex);
            parent.getChildren().add(left);
            NewickTreeNode right = new NewickTreeNode(0, 0, "", parent);
            reconstructTree(right, rows, depth + 1, rightTaxaList, taxaMapIndex);
            parent.getChildren().add(right);
        }
    }

    static class RowSum {
        String row;
        int freq;
        int n;
        double entropy;

        public RowSum(String row, int n) {
            this.row = row;
            this.n = n;
            this.freq = 0;
            for (int j = 0; j < row.length(); ++j) {
                if (row.charAt(j) == '1') {
                    ++freq;
                }
            }
            //熵值计算，熵值代表的是序列的混乱程度，序列的0和1数量越接近则混乱程度越高（最高的熵值是一半为0一半为1），全是0或全是1的熵值最低
            if (freq == 0 || freq == n) {
                entropy = 0;
            } else {
                double p1 = (double) freq / n;
                double p2 = 1.0 - p1;
                entropy = -p1 * Math.log(p1) - p2 * Math.log(p2);
            }
        }
    }
}
